from transformers import pipeline
import numpy as np
import scipy
# example of pipelien text-to-speech
# pipe=pipeline("text-to-speech", model="suno/bark-small", device='cuda')
# text="[clears throat] This is a test ...  and I just took a long pause."
# output=pipe(text)
model_path = "./bark-small/angelala00/bark-small"
pipe = pipeline("text-to-speech", model=model_path, device='cuda')
# text="[clears throat] This is a test ------—————...  and I just took a long pause."
text="今天天气不错 你知道吗 ？____------—————...  and I just took a long pause."
output=pipe(text)
# Step 3: 提取音频数据和采样率
audio = output["audio"]
sampling_rate = output["sampling_rate"]

# ✅ 关键修复：处理音频维度
audio = np.array(audio)          # 确保是 numpy 数组
audio = audio.squeeze()          # 去除所有长度为1的维度

# 确保是 1D 或 2D
if audio.ndim == 1:
    print(f"Audio is mono, shape: {audio.shape}")
elif audio.ndim == 2:
    if audio.shape[1] > 2:
        print(f"Trimming audio from {audio.shape[1]} channels to 2")
        audio = audio[:, :2]
    print(f"Audio is stereo, shape: {audio.shape}")
else:
    raise ValueError(f"Unsupported audio dimensions: {audio.shape}")

# 转为 float32
audio = audio.astype(np.float32)

# 保存
output_file = "output_audio.wav"
print(f"Saving audio to {output_file}...")
scipy.io.wavfile.write(output_file, rate=sampling_rate, data=audio)
print(f"✅ Audio saved to '{output_file}'")